package com.vf.cloud.rendering.common.util;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.LinkedList;
import java.util.List;

import org.dom4j.Document;
import org.dom4j.DocumentException;
import org.dom4j.Element;
import org.dom4j.io.SAXReader;

import com.jfinal.kit.StrKit;
import com.sun.jna.Platform;
import com.vf.cloud.rendering.common.constant.Cache;
import com.vf.cloud.rendering.common.constant.SoftPath;
import com.vf.cloud.rendering.common.vo.Gpu;

import lombok.extern.slf4j.Slf4j;

@Slf4j
public class GpuUtil {

	public static List<Gpu> getGPUs() {
		String s = sendGPUCmd();
		if (!StrKit.isBlank(s)) {
			return readGpuList(s);
		}
		return new LinkedList<Gpu>();
	}

	private static String sendGPUCmd() {
		Process process = null;
		try {
			if (Platform.isWindows()) {
				process = Runtime.getRuntime()
						.exec(String.format("%s/graphics/external/nvidia-smi -q -x", SoftPath.path));

			} else if (Platform.isLinux()) {
				String[] shell = { "/bin/bash", "-c", "nvidia-smi -q -x" };
				process = Runtime.getRuntime().exec(shell);
			}
			process.getOutputStream().close();
			BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
			StringBuffer stringBuffer = new StringBuffer();
			String line = "";
			while (null != (line = reader.readLine())) {
				stringBuffer.append(line + "\n");
			}
			return stringBuffer.toString();
		} catch (IOException e) {
			log.error("getGPU Error", e);
		}
		return null;
	}

	@SuppressWarnings("unchecked")
	private static List<Gpu> readGpuList(String s) {
		List<Gpu> list = new LinkedList<Gpu>();
		InputStream stream = new ByteArrayInputStream(s.getBytes(StandardCharsets.UTF_8));
		try {
			SAXReader reader = new SAXReader();
			reader.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);
			Document document = reader.read(stream);
			Element root = document.getRootElement();
			List<Element> elements = root.elements();
			Gpu gpuVo;
			for (Element element : elements) {
				if (StrKit.equals("timestamp", element.getQName().getName())) {

				} else if (StrKit.equals("driver_version", element.getQName().getName())) {

				} else if (StrKit.equals("cuda_version", element.getQName().getName())) {

				} else if (StrKit.equals("attached_gpus", element.getQName().getName())) {

				} else if (StrKit.equals("gpu", element.getQName().getName())) {
					List<Element> gpu = element.elements();
					if (gpu != null && gpu.size() > 0) {
						gpuVo = new Gpu();
						if (Cache.local != null)
							gpuVo.setMac(Cache.local.getMac());
						for (Element e : gpu) {
							if (StrKit.equals("product_name", e.getQName().getName())) {
								gpuVo.setName(e.getText());
							} else if (StrKit.equals("uuid", e.getQName().getName())) {
								gpuVo.setUuid(e.getText());
							} else if (StrKit.equals("board_id", e.getQName().getName())) {
								gpuVo.setBoardId(e.getText());
							} else if (StrKit.equals("vbios_version", e.getQName().getName())) {
								gpuVo.setVbiosVersion(e.getText());
							} else if (StrKit.equals("gpu_module_id", e.getQName().getName())) {
								gpuVo.setModuleId(Integer.parseInt(e.getText()));
							}

							else if (StrKit.equals("fb_memory_usage", e.getQName().getName())) {
								List<Element> fb_memory_usage = e.elements();
								for (Element f : fb_memory_usage) {
									if (StrKit.equals("total", f.getQName().getName())) {
										gpuVo.setMemoryTotal(Double.parseDouble(f.getText().replace("MiB","").trim())/1024);
									} else if (StrKit.equals("reserved", f.getQName().getName())) {
										gpuVo.setMemoryReserved(f.getText());
									} else if (StrKit.equals("used", f.getQName().getName())) {
										gpuVo.setMemoryUsed(f.getText());
									} else if (StrKit.equals("free", f.getQName().getName())) {
										gpuVo.setMemoryFree(f.getText());
									}
								}
							}

							else if (StrKit.equals("utilization", e.getQName().getName())) {
								List<Element> utilization = e.elements();
								for (Element f : utilization) {
									if (StrKit.equals("gpu_util", f.getQName().getName())) {
										gpuVo.setGpuUtil(f.getText());
									} else if (StrKit.equals("memory_util", f.getQName().getName())) {
										gpuVo.setMemoryUtil(f.getText());
									} else if (StrKit.equals("encoder_util", f.getQName().getName())) {
										gpuVo.setEncoderUtil(f.getText());
									} else if (StrKit.equals("decoder_util", f.getQName().getName())) {
										gpuVo.setDecoderUtil(f.getText());
									}
								}
							}
						}
						list.add(gpuVo);
					}
				}
			}
		} catch (DocumentException e) {
			e.printStackTrace();
		} catch (SecurityException e) {
			e.printStackTrace();
		} catch (Exception e) {
			e.printStackTrace();
		} finally {
			try {
				stream.close();
			} catch (IOException e) {
				e.printStackTrace();
			}
		}
		return list;
	}

}
